# -*- coding: utf-8 -*-
import json
import requests
from configparser import ConfigParser
# import config_utils

class ChatCompletion:
    def __init__(self, config):
        """
        初始化 ChatCompletion 类
        :param config: 配置对象
        """
        self.config = config
        self.api_url = self.config.get("api_setting.api_url")
        self.api_token = self.config.get("api_setting.api_token")
        self.model = self.config.get("api_setting.model")
        self.headers = {
            "Authorization": f"Bearer {self.api_token}",
            "Content-Type": "application/json"
        }

    def generate_payload(self, messages, temperature=0.7):
        """
        生成请求的 payload
        :param messages: 对话消息
        :param temperature: 温度参数
        :return: 请求的 payload
        """
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": messages}],
            "temperature": temperature,
            "stream": True
        }
        return payload

    def send_request(self, payload):
        """
        发送请求并返回响应
        :param payload: 请求的 payload
        :return: API 响应
        """
        return requests.post(self.api_url, json=payload, headers=self.headers, stream=True)

    def get_chat_completion(self, messages):
        """
        获取聊天补全结果
        :param messages: 对话消息
        :return: 生成器，逐块返回内容
        """
        try:
            payload = self.generate_payload(messages)
            response = self.send_request(payload)
            if response.status_code == 200:
                for line in response.iter_lines():
                    if line:
                        content = self.extract_content(line)
                        if content:
                            yield content
            else:
                print(f"Request failed with status code: {response.status_code}")
                print(f"Response content: {response.text}")
        except requests.RequestException as e:
            print(f"Request failed: {e}")

    def extract_content(self, line):
        """
        提取 content 字段
        :param line: SSE 数据行
        :return: 提取的 content 或 None
        """
        # 解码字节数据
        line_str = line.decode('utf-8')

        # 检查是否结束
        if line_str.strip() == "data: [DONE]":
            return None

        # 去掉 "data: " 前缀
        if line_str.startswith("data: "):
            json_str = line_str[len("data: "):].strip()

            # 解析 JSON
            try:
                data = json.loads(json_str)
                # 提取 content
                if "choices" in data and len(data["choices"]) > 0:
                    content = data["choices"][0]["delta"].get("content")
                    return self.remove_symbols(content) if content else ""
            except json.JSONDecodeError as e:
                print(f"Failed to parse JSON: {e}")
        return None

    def remove_symbols(self, text):
        """
        去掉文本中的 #、*、[、]、\\ 符号
        :param text: 原始文本
        :return: 去掉符号后的文本
        """
        symbols = ["#", "*", "[", "]", "\\"]
        for symbol in symbols:
            text = text.replace(symbol, "")
        return text

# 使用示例
if __name__ == "__main__":
    path = "../config.yaml"
    import config_utils
    config = config_utils.ConfigManager(file_path=path)

    chat_completion = ChatCompletion(config)
    messages = "10是质数吗？"

    for content in chat_completion.get_chat_completion(messages=messages):
        print(content, end="", flush=True)


